import numpy as np
import matplotlib

matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt
import pandas as pd
from math import isclose
import random
from collections import Counter
import numpy.random as npr
import numpy.linalg as linalg

np.set_printoptions(precision=4, suppress=True)


# 对于有限状态的不可约马尔可夫链，平稳分布存在且唯一
# 平稳分布pi=pi * p  即与转移矩阵相乘后保持不变
# 对于不可约、非周期性的马尔科夫链，无论从哪个初始状态开始，经验分布都会收敛到平稳分布：
# pi = [a,b,c]
# a = 0.6a + 0.5c
# b = 0.4a + 0.5c
# c=b
# a+b+c=1
# 解得：a = 5/13 ≈ 0.3846, b = 4/13 ≈ 0.3077, c = 4/13 ≈ 0.3077


def simulate_chain(P, start, n_steps):
    """
    Simulate Markov chain with transition matrix P starting from state index start
    for n_steps.
    Returns the list of visted states including the start state.
    :param P:
    :param start:
    :param n_steps:
    :return:
    """
    states = [start]
    cur = start
    for _ in range(n_steps):
        cur = npr.choice(a=len(P), p=P[cur])
        states.append(cur)

    return states


print("Section1: Example  - 'Good irreducible chain")
'''
这个转移矩阵描述了一个不可约马尔可夫链：
状态0 → 状态0或状态1
状态1 → 状态2
状态2 → 状态0或状态1
所有状态都是互相可达的，因此这个链是不可约的。
'''
P1 = np.array([
    [0.6, 0.4, 0.0],  # 状态0: 60%留在0, 40%转到1
    [0.0, 0.0, 1.0],  # 状态1: 100%转到2
    [0.5, 0.5, 0.0]  # 状态2: 50%转到0, 50%转到1
])
print("Transition matrix P1:\n", P1)


# simulate many steps and show empirical distribution convergence
def empirical_distribution(P, start, steps):
    '''
    通过长时间模拟马尔可夫链
    统计各个状态在轨迹中出现的频率
    用频率来近似概率分布
    :param P:
    :param start:
    :param steps:
    :return:
    '''
    traj = simulate_chain(P, start, steps)
    counts = Counter(traj[-10000:]) if steps > 10000 else Counter(traj)
    vec = np.array([counts[i] / sum(counts.values()) for i in range(len(P))])
    return vec


emp = empirical_distribution(P1, start=0, steps=20000)
print("Empirical distribution (last portion of trajectory):", emp)


# compute stationary via power method
def stationary_power_method(P, tol=1e-10, max_iter=10000):
    n = P.shape[0] # 状态数量
    v = np.ones(n) / n  # 初始化均匀分布向量 [1/n, 1/n, ..., 1/n]
    for _ in range(max_iter):
        v_next = v.dot(P)
        # ||v_next - v||₁ = Σ|v_next[i] - v[i]|
        # 使用 L1 范数是因为它对于概率分布有自然的解释：两个概率分布之间的总变分距离。
        if linalg.norm(v_next - v, ord=1) < tol: # 检查收敛
            return v_next
        v = v_next
    return v

# 初始化：v = [1/3, 1/3, 1/3]（对于3状态系统）
# 第一次迭代：v₁ = v₀ × P
# 第二次迭代：v₂ = v₁ × P = v₀ × P²
# 第k次迭代：vₖ = v₀ × Pᵏ
#随着 k → ∞，vₖ 收敛到平稳分布 π。

s1 = stationary_power_method(P1)
print("Stationary distribution via power method:", s1)
